package cn.iocoder.yudao.framework.handlermapping;

import cn.iocoder.yudao.framework.common.exception.ServiceException;
import cn.iocoder.yudao.framework.message.ClientRequest;
import cn.iocoder.yudao.framework.message.ServerResponse;
import cn.iocoder.yudao.framework.server.websocket.security.SecurityFrameworkUtils;
import cn.iocoder.yudao.framework.server.websocket.security.WebSocketUser;
import cn.iocoder.yudao.module.infra.util.TraceIdUtil;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author 王广帅
 * @date 2021年01月26日 4:26 下午
 */
public class NettyWsChannelContext {

    public static final AttributeKey<WsSessionContext> WEB_SOCKET_USER_BIND_ATTRIBUTE_KEY = AttributeKey.valueOf(
            "web_socket_channel_user_attribute_key");
    public static final AttributeKey<String> TRACE_ID = AttributeKey.valueOf(TraceIdUtil.TID);

    private static final Logger logger = LoggerFactory.getLogger(NettyWsChannelContext.class);
    private ClientRequest clientRequest;
    private ChannelHandlerContext nettyCtx;

    public NettyWsChannelContext(ChannelHandlerContext nettyCtx, ClientRequest clientRequest) {
        this.clientRequest = clientRequest;
        this.nettyCtx = nettyCtx;
    }

    public ChannelHandlerContext getNettyCtx() {
        return nettyCtx;
    }

    public WebSocketUser getUser() {
        final Channel channel = nettyCtx.channel();
        final Attribute<Object> attr = channel.attr(AttributeKey.valueOf(SecurityFrameworkUtils.WEB_SOCKET_USER_ATTRIBUTE_KEY));
        final WebSocketUser webSocketUser = (WebSocketUser) attr.get();
        return webSocketUser;
    }

    public void writeAndFlush(Object msg) {
        ServerResponse response = new ServerResponse();
        response.setUri(this.clientRequest.getUri());
        response.setData(msg);
        if (nettyCtx.channel().isWritable()) {
            nettyCtx.channel().attr(TRACE_ID).set(TraceIdUtil.getTraceId());
            nettyCtx.writeAndFlush(response);
        } else {
            logger.warn("数据未发送成功，被拦截了, msg = {}", msg);
        }
    }

    public void writeAndFlush(Throwable exception) {

        if (exception instanceof ServiceException) {
            logger.warn("服务器异常", exception);
            ServiceException logicException = (ServiceException) exception;
            ServerResponse serverError = new ServerResponse();
            serverError.setCode(logicException.getCode());
            serverError.setMessage(logicException.getMessage());
            writeAndFlush(serverError);
        } else {
            ServerResponse serverError = new ServerResponse();
            serverError.setCode(-1);
            serverError.setMessage("系统异常:" + exception.getMessage());
            logger.error("系统异常", exception);
            writeAndFlush(serverError);
        }
    }

    public WsSessionContext getSession() {
        Attribute<WsSessionContext> attr = this.nettyCtx.channel().attr(WEB_SOCKET_USER_BIND_ATTRIBUTE_KEY);
        if (attr != null) {
            return attr.get();
        }
        return null;

    }

}
